import pickle

import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm

from models.LeNet import LeNet
from models.MobileNet import MobileNetCIFAR
from utils.TrainingMonitor import TrainingMonitor


def load_cifar_batch(file_path):
    """修正版数据加载函数（网页2/5/7）"""
    with open(file_path, 'rb') as f:
        batch = pickle.load(f, encoding='latin1')

    # 数据格式转换（NHWC -> NCHW）
    data = batch['data'].reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)  # 转换为NHWC
    data = data.transpose(0, 3, 1, 2)  # 转换为PyTorch需要的NCHW格式
    return data.astype(np.float32), np.array(batch['labels'])


# 加载训练集（网页7）
train_data, train_labels = [], []
for i in range(1, 6):
    data, labels = load_cifar_batch(f'cifar-10-batches-py/data_batch_{i}')
    train_data.append(data)
    train_labels.append(labels)
train_data = np.concatenate(train_data)
train_labels = np.concatenate(train_labels)

# 加载测试集（网页3）
test_data, test_labels = load_cifar_batch('cifar-10-batches-py/test_batch')

# 数据标准化（网页7）
mean = np.array([0.4914, 0.4822, 0.4465]).reshape(1, 3, 1, 1)
std = np.array([0.247, 0.243, 0.261]).reshape(1, 3, 1, 1)
train_data = (train_data / 255.0 - mean) / std
test_data = (test_data / 255.0 - mean) / std

# 创建DataLoader（网页1）
train_dataset = torch.utils.data.TensorDataset(
    torch.tensor(train_data, dtype=torch.float32),  # 确保创建时指定float32
    torch.tensor(train_labels, dtype=torch.long)
)
test_dataset = TensorDataset(
    torch.tensor(test_data, dtype=torch.float32),  # 确保创建时指定float32
    torch.tensor(test_labels, dtype=torch.long)
)

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=128)


def evaluate_loss(model, loader, device, criterion):
    model.eval()
    total_loss = 0.0
    with torch.no_grad():
        for inputs, labels in loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            total_loss += loss.item()
    return total_loss / len(loader)


def evaluate(model: nn.Module, loader, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in loader:
            # 强制四维输入（网页1/7）
            if inputs.dim() != 4:
                inputs = inputs.view(-1, 3, 32, 32).contiguous()

            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return correct / total


def train_model(model: nn.Module, epochs=20):
    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    model = model.to(device)
    model.shape_debug = True  # 开启形状调试

    # 验证初始批次形状
    test_batch = next(iter(train_loader))[0].to(device).float()
    print("验证输入形状:", test_batch.shape)  # 应显示[128,3,32,32]
    _ = model(test_batch)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    monitor = TrainingMonitor(model.__class__.__name__)
    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0

        pbar = tqdm(train_loader, desc=f'Epoch {epoch + 1}')

        for inputs, labels in pbar:
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            # 统计训练指标
            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

            pbar.set_postfix({'loss': loss.item()})

        # 计算训练准确率
        train_acc = 100. * correct / total

        # 验证阶段
        val_acc = evaluate(model, test_loader, device) * 100
        val_loss = evaluate_loss(model, test_loader, device, criterion)

        # 记录指标
        monitor.record(epoch + 1, running_loss / len(train_loader), val_loss, train_acc, val_acc)
        print(f'\rEpoch {epoch + 1} | Loss: {running_loss / len(train_loader):.4f} | Val Acc: {val_acc:.2f}%', flush=True)

        # 生成训练曲线
    monitor.plot_curves()
    return model


if __name__ == '__main__':
    # 示例：训练MobileNet
    for model in [LeNet]:

        trained_model = train_model(model(), epochs=20)

        # 保存模型
        torch.save(trained_model.state_dict(), f"weights/cifar10_{model.__name__}.pth")
